"""
特征提取模块

基于reference/1-2-Multi-dimensional_feature_extraction.py改进
实现多维特征提取功能
"""

import numpy as np
import pandas as pd
import scipy.signal as signal
import pywt
from scipy.stats import entropy
from scipy.spatial.distance import pdist, squareform
import warnings
warnings.filterwarnings('ignore')


class FeatureExtractor:
    """特征提取器 - 基于reference/1-2改进"""
    
    def __init__(self, target_fs=12000):
        self.target_fs = target_fs
        self.feature_names = []
        
    def extract_all_features(self, signal_data, rpm, bearing_type='SKF6205'):
        """提取所有类型的特征"""
        features = {}
        
        # 时域特征
        time_features = self._extract_time_domain_features(signal_data)
        features.update(time_features)
        
        # 频域特征
        freq_features = self._extract_frequency_domain_features(signal_data, self.target_fs)
        features.update(freq_features)
        
        # 时频域特征
        time_freq_features = self._extract_time_frequency_features(signal_data, self.target_fs)
        features.update(time_freq_features)
        
        # 轴承故障特征频率
        fault_freq_features = self._extract_fault_frequency_features(signal_data, self.target_fs, rpm, bearing_type)
        features.update(fault_freq_features)
        
        # 非线性特征
        nonlinear_features = self._extract_nonlinear_features(signal_data)
        features.update(nonlinear_features)
        
        return features
    
    def _extract_time_domain_features(self, signal):
        """提取时域特征"""
        features = {}
        
        # 基本统计特征
        features['mean'] = np.mean(signal)
        features['std'] = np.std(signal)
        features['max'] = np.max(signal)
        features['min'] = np.min(signal)
        features['peak_to_peak'] = np.ptp(signal)
        features['rms'] = np.sqrt(np.mean(signal**2))
        
        # 形状特征
        features['skewness'] = self._calculate_skewness(signal)
        features['kurtosis'] = self._calculate_kurtosis(signal)
        features['crest_factor'] = np.max(np.abs(signal)) / features['rms'] if features['rms'] != 0 else 0
        features['shape_factor'] = features['rms'] / np.mean(np.abs(signal)) if np.mean(np.abs(signal)) != 0 else 0
        features['impulse_factor'] = np.max(np.abs(signal)) / np.mean(np.abs(signal)) if np.mean(np.abs(signal)) != 0 else 0
        features['clearance_factor'] = np.max(np.abs(signal)) / (np.mean(np.sqrt(np.abs(signal)))**2) if np.mean(np.sqrt(np.abs(signal))) != 0 else 0
        
        return features
    
    def _extract_frequency_domain_features(self, signal, fs):
        """提取频域特征"""
        features = {}
        
        # 计算FFT
        fft = np.fft.fft(signal)
        freqs = np.fft.fftfreq(len(signal), 1/fs)
        magnitude = np.abs(fft)
        
        # 只取正频率部分
        positive_freqs = freqs[:len(freqs)//2]
        positive_magnitude = magnitude[:len(magnitude)//2]
        
        # 基本频域特征
        features['freq_mean'] = np.mean(positive_magnitude)
        features['freq_std'] = np.std(positive_magnitude)
        features['freq_max'] = np.max(positive_magnitude)
        features['freq_centroid'] = np.sum(positive_freqs * positive_magnitude) / np.sum(positive_magnitude) if np.sum(positive_magnitude) != 0 else 0
        
        # 频带能量特征
        total_energy = np.sum(positive_magnitude**2)
        if total_energy > 0:
            # 低频带 (0-1/3 fs)
            low_freq_mask = positive_freqs <= fs/3
            features['band_1_energy'] = np.sum(positive_magnitude[low_freq_mask]**2) / total_energy
            
            # 中频带 (1/3 fs - 2/3 fs)
            mid_freq_mask = (positive_freqs > fs/3) & (positive_freqs <= 2*fs/3)
            features['band_2_energy'] = np.sum(positive_magnitude[mid_freq_mask]**2) / total_energy
            
            # 高频带 (2/3 fs - fs)
            high_freq_mask = positive_freqs > 2*fs/3
            features['band_3_energy'] = np.sum(positive_magnitude[high_freq_mask]**2) / total_energy
        else:
            features['band_1_energy'] = 0
            features['band_2_energy'] = 0
            features['band_3_energy'] = 0
        
        return features
    
    def _extract_time_frequency_features(self, signal, fs):
        """提取时频域特征"""
        features = {}
        
        # 小波变换特征
        try:
            # 选择小波基
            wavelet = 'db4'
            # 计算小波分解层数
            max_level = min(4, int(np.log2(len(signal))) - 1)
            
            coeffs = pywt.wavedec(signal, wavelet, level=max_level)
            
            # 各层小波系数特征
            for i, coeff in enumerate(coeffs):
                if len(coeff) > 0:
                    features[f'wavelet_{i}_energy'] = np.sum(coeff**2)
                    features[f'wavelet_{i}_std'] = np.std(coeff)
        except:
            # 如果小波变换失败，使用默认值
            for i in range(5):
                features[f'wavelet_{i}_energy'] = 0
                features[f'wavelet_{i}_std'] = 0
        
        # STFT特征
        try:
            # 动态调整STFT参数
            nperseg = min(256, len(signal) // 4)
            noverlap = nperseg // 2
            
            f, t, Zxx = signal.stft(signal, fs, nperseg=nperseg, noverlap=noverlap)
            
            # STFT统计特征
            features['stft_mean'] = np.mean(np.abs(Zxx))
            features['stft_std'] = np.std(np.abs(Zxx))
        except:
            features['stft_mean'] = 0
            features['stft_std'] = 0
        
        return features
    
    def _extract_fault_frequency_features(self, signal, fs, rpm, bearing_type):
        """提取轴承故障特征频率"""
        features = {}
        
        # 轴承参数
        bearing_params = {
            'SKF6205': {'n': 9, 'd': 0.3126, 'D': 1.537},
            'SKF6203': {'n': 9, 'd': 0.2656, 'D': 1.122}
        }
        
        if bearing_type not in bearing_params:
            bearing_type = 'SKF6205'
        
        n = bearing_params[bearing_type]['n']
        d = bearing_params[bearing_type]['d']
        D = bearing_params[bearing_type]['D']
        
        # 计算故障特征频率
        fr = rpm / 60.0  # 转频
        
        # 外圈故障频率 (BPFO)
        bpfo = n * fr * (1 - d/D * np.cos(np.pi/2)) / 2
        
        # 内圈故障频率 (BPFI)
        bpfi = n * fr * (1 + d/D * np.cos(np.pi/2)) / 2
        
        # 滚动体故障频率 (BSF)
        bsf = fr * D * (1 - (d/D * np.cos(np.pi/2))**2) / (2 * d)
        
        # 保持架故障频率 (FTF)
        ftf = fr * (1 - d/D * np.cos(np.pi/2)) / 2
        
        # 计算各故障频率的谐波能量
        fault_freqs = [bpfo, bpfi, bsf, ftf]
        fault_names = ['bpfo', 'bpfi', 'bsf', 'ftf']
        
        for i, (freq, name) in enumerate(zip(fault_freqs, fault_names)):
            for harmonic in range(1, 4):  # 1-3次谐波
                target_freq = freq * harmonic
                if target_freq < fs/2:  # 确保在奈奎斯特频率内
                    # 计算目标频率附近的能量
                    energy = self._calculate_frequency_energy(signal, fs, target_freq, bandwidth=10)
                    features[f'{name}_harmonic_{harmonic}_energy'] = energy
                else:
                    features[f'{name}_harmonic_{harmonic}_energy'] = 0
        
        return features
    
    def _extract_nonlinear_features(self, signal):
        """提取非线性特征"""
        features = {}
        
        # 近似熵
        features['approx_entropy'] = self._calculate_approximate_entropy(signal)
        
        # 样本熵
        features['sample_entropy'] = self._calculate_sample_entropy(signal)
        
        # 排列熵
        features['permutation_entropy'] = self._calculate_permutation_entropy(signal)
        
        return features
    
    def _calculate_skewness(self, signal):
        """计算偏度"""
        mean = np.mean(signal)
        std = np.std(signal)
        if std == 0:
            return 0
        return np.mean(((signal - mean) / std) ** 3)
    
    def _calculate_kurtosis(self, signal):
        """计算峰度"""
        mean = np.mean(signal)
        std = np.std(signal)
        if std == 0:
            return 0
        return np.mean(((signal - mean) / std) ** 4)
    
    def _calculate_frequency_energy(self, signal, fs, target_freq, bandwidth=10):
        """计算目标频率附近的能量"""
        # 计算FFT
        fft = np.fft.fft(signal)
        freqs = np.fft.fftfreq(len(signal), 1/fs)
        
        # 只取正频率部分
        positive_freqs = freqs[:len(freqs)//2]
        positive_fft = fft[:len(fft)//2]
        
        # 找到目标频率附近的索引
        freq_mask = (positive_freqs >= target_freq - bandwidth/2) & (positive_freqs <= target_freq + bandwidth/2)
        
        # 计算能量
        energy = np.sum(np.abs(positive_fft[freq_mask])**2)
        
        return energy
    
    def _calculate_approximate_entropy(self, signal, m=2, r=0.2):
        """计算近似熵"""
        if len(signal) < 100:
            return 0.5  # 简化实现
        
        # 对长信号进行下采样
        if len(signal) > 1000:
            step = len(signal) // 1000
            signal = signal[::step]
        
        N = len(signal)
        if N < m + 1:
            return 0.5
        
        # 简化实现，避免复杂计算
        try:
            # 计算阈值
            r = r * np.std(signal)
            
            # 简化的近似熵计算
            def _maxdist(xi, xj, m):
                return max([abs(ua - va) for ua, va in zip(xi, xj)])
            
            def _approximate_entropy(m):
                C = np.zeros(N - m + 1)
                # 限制计算范围以提高性能
                max_iter = min(100, N - m + 1)
                for i in range(max_iter):
                    template_i = signal[i:i + m]
                    count = 0
                    for j in range(max_iter):
                        template_j = signal[j:j + m]
                        if _maxdist(template_i, template_j, m) <= r:
                            count += 1.0
                    C[i] = count
                
                phi = np.mean(np.log(C / max_iter + 1e-10))
                return phi
            
            return _approximate_entropy(m) - _approximate_entropy(m + 1)
        except:
            return 0.5
    
    def _calculate_sample_entropy(self, signal, m=2, r=0.2):
        """计算样本熵"""
        if len(signal) < 100:
            return 1.0  # 简化实现
        
        # 对长信号进行下采样
        if len(signal) > 1000:
            step = len(signal) // 1000
            signal = signal[::step]
        
        N = len(signal)
        if N < m + 1:
            return 1.0
        
        # 简化实现，避免复杂计算
        try:
            # 计算阈值
            r = r * np.std(signal)
            
            # 简化的样本熵计算
            def _maxdist(xi, xj, m):
                return max([abs(ua - va) for ua, va in zip(xi, xj)])
            
            def _sample_entropy(m):
                B = 0.0
                A = 0.0
                
                # 限制计算范围以提高性能
                max_iter = min(50, N - m)
                for i in range(max_iter):
                    template_i = signal[i:i + m]
                    for j in range(max_iter):
                        if i != j:
                            template_j = signal[j:j + m]
                            if _maxdist(template_i, template_j, m) <= r:
                                B += 1.0
                                
                                # 检查m+1维模板
                                if i + m < N and j + m < N:
                                    template_i_m1 = signal[i:i + m + 1]
                                    template_j_m1 = signal[j:j + m + 1]
                                    if _maxdist(template_i_m1, template_j_m1, m + 1) <= r:
                                        A += 1.0
                
                if B == 0:
                    return 0.0
                
                return -np.log(A / B + 1e-10)
            
            return _sample_entropy(m)
        except:
            return 1.0
    
    def _calculate_permutation_entropy(self, signal, m=3, delay=1):
        """计算排列熵"""
        if len(signal) < 50:
            return 0.0  # 简化实现
        
        # 对长信号进行下采样
        if len(signal) > 500:
            step = len(signal) // 500
            signal = signal[::step]
        
        N = len(signal)
        if N < m:
            return 0.0
        
        # 简化实现，避免复杂计算
        try:
            # 创建延迟向量
            patterns = []
            max_patterns = min(200, N - (m - 1) * delay)
            for i in range(max_patterns):
                pattern = signal[i:i + m * delay:delay]
                if len(pattern) == m:
                    patterns.append(pattern)
            
            if len(patterns) < 10:
                return 0.0
            
            # 计算排列模式
            permutation_patterns = []
            for pattern in patterns:
                # 获取排序索引
                sorted_indices = np.argsort(pattern)
                permutation_patterns.append(tuple(sorted_indices))
            
            # 计算排列熵
            unique_patterns, counts = np.unique(permutation_patterns, return_counts=True)
            probabilities = counts / len(permutation_patterns)
            
            # 计算熵
            entropy_val = -np.sum(probabilities * np.log2(probabilities + 1e-10))
            
            return entropy_val
        except:
            return 0.0
